#include "mainwindow.h"
#include <QVBoxLayout>
#include <QHBoxLayout>
#include <QWidget>
#include <QFileDialog>
#include <QMessageBox>
#include <QDir>
#include <QDebug>
#include <QSpinBox>
#include <QLabel>

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
{
    setupUi();
    initializeKeywords();
}

MainWindow::~MainWindow() {}

void MainWindow::setupUi() {
    // 创建中心部件和主布局
    QWidget* centralWidget = new QWidget(this);
    QVBoxLayout* mainLayout = new QVBoxLayout(centralWidget);
    
    // JSON处理组
    QGroupBox* jsonGroup = new QGroupBox("JSON处理", this);
    QVBoxLayout* jsonLayout = new QVBoxLayout(jsonGroup);
    
    // JSON文件选择
    QHBoxLayout* jsonFileSelectLayout = new QHBoxLayout();
    jsonFilePathEdit = new QLineEdit(this);
    jsonFilePathEdit->setPlaceholderText("选择JSON文件路径");
    selectJsonFileBtn = new QPushButton("浏览", this);
    jsonFileSelectLayout->addWidget(jsonFilePathEdit);
    jsonFileSelectLayout->addWidget(selectJsonFileBtn);
    
    // JSON关键词输入
    jsonKeywordsEdit = new QLineEdit(this);
    jsonKeywordsEdit->setPlaceholderText("输入关键词，用逗号分隔（例如：快乐,开心）");
    
    processJsonBtn = new QPushButton("整理JSON生成TXT", this);
    
    // 添加到JSON组
    jsonLayout->addLayout(jsonFileSelectLayout);
    jsonLayout->addWidget(jsonKeywordsEdit);
    jsonLayout->addWidget(processJsonBtn);
    
    // TXT处理组
    QGroupBox* txtGroup = new QGroupBox("TXT处理", this);
    QVBoxLayout* txtLayout = new QVBoxLayout(txtGroup);
    
    // TXT文件选择
    QHBoxLayout* txtFileSelectLayout = new QHBoxLayout();
    txtFilePathEdit = new QLineEdit(this);
    txtFilePathEdit->setPlaceholderText("选择TXT文件路径");
    selectTxtFileBtn = new QPushButton("浏览", this);
    txtFileSelectLayout->addWidget(txtFilePathEdit);
    txtFileSelectLayout->addWidget(selectTxtFileBtn);
    
    // TXT长度限制设置
    QHBoxLayout* lengthLimitLayout = new QHBoxLayout();
    QLabel* minLengthLabel = new QLabel("最小长度:", this);
    minLengthSpinBox = new QSpinBox(this);
    minLengthSpinBox->setRange(1, 1000);
    minLengthSpinBox->setValue(10);
    minLengthSpinBox->setToolTip("设置分割文本的最小长度");
    
    QLabel* maxLengthLabel = new QLabel("最大长度:", this);
    maxLengthSpinBox = new QSpinBox(this);
    maxLengthSpinBox->setRange(10, 4096);
    maxLengthSpinBox->setValue(512);
    maxLengthSpinBox->setToolTip("设置分割文本的最大长度");
    
    lengthLimitLayout->addWidget(minLengthLabel);
    lengthLimitLayout->addWidget(minLengthSpinBox);
    lengthLimitLayout->addSpacing(20);
    lengthLimitLayout->addWidget(maxLengthLabel);
    lengthLimitLayout->addWidget(maxLengthSpinBox);
    
    // TXT关键词输入
    txtKeywordsEdit = new QLineEdit(this);
    txtKeywordsEdit->setPlaceholderText("输入TXT关键词，用逗号分隔（例如：句子1,句子2）");
    
    processTxtBtn = new QPushButton("整理TXT生成SRT", this);
    
    // 添加到TXT组
    txtLayout->addLayout(txtFileSelectLayout);
    txtLayout->addLayout(lengthLimitLayout);
    txtLayout->addWidget(txtKeywordsEdit);
    txtLayout->addWidget(processTxtBtn);
    
    // SRT处理组
    QGroupBox* srtGroup = new QGroupBox("SRT处理", this);
    QVBoxLayout* srtLayout = new QVBoxLayout(srtGroup);
    
    // 未翻译SRT文件选择
    QHBoxLayout* untranslatedSrtLayout = new QHBoxLayout();
    untranslatedSrtPathEdit = new QLineEdit(this);
    untranslatedSrtPathEdit->setPlaceholderText("选择未翻译SRT文件路径");
    selectUntranslatedSrtBtn = new QPushButton("浏览", this);
    untranslatedSrtLayout->addWidget(untranslatedSrtPathEdit);
    untranslatedSrtLayout->addWidget(selectUntranslatedSrtBtn);

    // 已翻译SRT文件选择
    QHBoxLayout* translatedSrtLayout = new QHBoxLayout();
    translatedSrtPathEdit = new QLineEdit(this);
    translatedSrtPathEdit->setPlaceholderText("选择已翻译SRT文件路径");
    selectTranslatedSrtBtn = new QPushButton("浏览", this);
    translatedSrtLayout->addWidget(translatedSrtPathEdit);
    translatedSrtLayout->addWidget(selectTranslatedSrtBtn);

    processSrtBtn = new QPushButton("整理SRT生成JSON", this);
    
    // 添加到SRT组
    srtLayout->addLayout(untranslatedSrtLayout);
    srtLayout->addLayout(translatedSrtLayout);
    srtLayout->addWidget(processSrtBtn);
    
    // TXT转JSON组
    QGroupBox* txtToJsonGroup = new QGroupBox("TXT转JSON", this);
    QVBoxLayout* txtToJsonLayout = new QVBoxLayout(txtToJsonGroup);

    // TXT文件选择
    QHBoxLayout* txtToJsonFileSelectLayout = new QHBoxLayout();
    txtToJsonInputEdit = new QLineEdit(this);
    txtToJsonInputEdit->setPlaceholderText("选择要转换为JSON的TXT文件");
    selectTxtToJsonInputBtn = new QPushButton("浏览", this);
    txtToJsonFileSelectLayout->addWidget(txtToJsonInputEdit);
    txtToJsonFileSelectLayout->addWidget(selectTxtToJsonInputBtn);

    processTxtToJsonBtn = new QPushButton("转换为JSON", this);

    // 添加到��局
    txtToJsonLayout->addLayout(txtToJsonFileSelectLayout);
    txtToJsonLayout->addWidget(processTxtToJsonBtn);

    // JSON转SRT组
    QGroupBox* jsonToSrtGroup = new QGroupBox("JSON转SRT", this);
    QVBoxLayout* jsonToSrtLayout = new QVBoxLayout(jsonToSrtGroup);

    // JSON文件选择
    QHBoxLayout* jsonToSrtFileSelectLayout = new QHBoxLayout();
    jsonToSrtInputEdit = new QLineEdit(this);
    jsonToSrtInputEdit->setPlaceholderText("选择要转换为SRT的JSON文件");
    selectJsonToSrtInputBtn = new QPushButton("浏览", this);
    jsonToSrtFileSelectLayout->addWidget(jsonToSrtInputEdit);
    jsonToSrtFileSelectLayout->addWidget(selectJsonToSrtInputBtn);

    processJsonToSrtBtn = new QPushButton("转换为SRT", this);

    // 添加到布局
    jsonToSrtLayout->addLayout(jsonToSrtFileSelectLayout);
    jsonToSrtLayout->addWidget(processJsonToSrtBtn);

    // 将所有组添加到主布局
    mainLayout->addWidget(jsonGroup);
    mainLayout->addWidget(txtGroup);
    mainLayout->addWidget(srtGroup);
    mainLayout->addWidget(txtToJsonGroup);
    mainLayout->addWidget(jsonToSrtGroup);
    
    // 设置中心部件
    setCentralWidget(centralWidget);
    
    // 连接信号和槽
    connect(selectJsonFileBtn, &QPushButton::clicked, this, &MainWindow::selectJsonFile);
    connect(processJsonBtn, &QPushButton::clicked, this, &MainWindow::processJson);
    connect(selectTxtFileBtn, &QPushButton::clicked, this, &MainWindow::selectTxtFile);
    connect(processTxtBtn, &QPushButton::clicked, this, &MainWindow::processTxt);
    connect(selectUntranslatedSrtBtn, &QPushButton::clicked, this, &MainWindow::selectUntranslatedSrt);
    connect(selectTranslatedSrtBtn, &QPushButton::clicked, this, &MainWindow::selectTranslatedSrt);
    connect(processSrtBtn, &QPushButton::clicked, this, &MainWindow::processSrt);
    connect(selectTxtToJsonInputBtn, &QPushButton::clicked, this, &MainWindow::selectTxtToJsonInput);
    connect(processTxtToJsonBtn, &QPushButton::clicked, this, &MainWindow::processTxtToJson);
    connect(selectJsonToSrtInputBtn, &QPushButton::clicked, this, &MainWindow::selectJsonToSrtInput);
    connect(processJsonToSrtBtn, &QPushButton::clicked, this, &MainWindow::processJsonToSrt);
    
    // 设置窗口属性
    setWindowTitle("数据集工具");
    resize(800, 600);  // 调整窗口大小以适应新布局
}

void MainWindow::selectJsonFile() {
    QString filePath = QFileDialog::getOpenFileName(
        this,
        "选择JSON文件",
        QString(),
        "JSON文件 (*.json);;所有文件 (*.*)"
    );
    
    if (!filePath.isEmpty()) {
        jsonFilePathEdit->setText(filePath);
    }
}

void MainWindow::processJson() {
    QString inputPath = jsonFilePathEdit->text();
    if (inputPath.isEmpty()) {
        QMessageBox::warning(this, "警告", "请先选择JSON文件");
        return;
    }

    // 获取JSON关键词列表
    QStringList keywords = getKeywordsList(jsonKeywordsEdit->text());
    
    // 构建输出路径
    QString outputDir = QDir::currentPath() + "/output";
    QString outputPath = outputDir + "/output.txt";

    // 保存关键词并处理JSON文件
    saveKeywords(keywords, jsonKeywordsFileName);
    
    if (jsonProcessor.processJsonToTxt(inputPath, outputPath, keywords)) {
        QMessageBox::information(this, "成功", "文件处理完成！\n保存路径: " + outputPath);
    } else {
        QMessageBox::critical(this, "错误", "文件处理失败！");
    }
}

void MainWindow::selectTxtFile() {
    QString filePath = QFileDialog::getOpenFileName(
        this,
        "选择TXT文件",
        QString(),
        "TXT文件 (*.txt);;所有文件 (*.*)"
    );
    
    if (!filePath.isEmpty()) {
        txtFilePathEdit->setText(filePath);
    }
}

void MainWindow::processTxt() {
    QString inputPath = txtFilePathEdit->text();
    if (inputPath.isEmpty()) {
        QMessageBox::warning(this, "警告", "请先选择TXT文件");
        return;
    }

    // 获取长度限制设置
    int minLength = minLengthSpinBox->value();
    int maxLength = maxLengthSpinBox->value();
    
    // 验证长度设置
    if (minLength >= maxLength) {
        QMessageBox::warning(this, "警告", "最小长度必须小于最大长度");
        return;
    }

    // 设置处理器的长度限制
    txtProcessor.setLengthLimits(minLength, maxLength);

    // 获取TXT关键词列表
    QStringList keywords = getKeywordsList(txtKeywordsEdit->text());
    
    // 构建输出路径
    QString outputDir = QDir::currentPath() + "/output";
    QString outputPath = outputDir + "/output.srt";

    // 保存关键词并处理TXT文件
    saveKeywords(keywords, txtKeywordsFileName);
    
    if (txtProcessor.processTextToSrt(inputPath, outputPath, keywords)) {
        QMessageBox::information(this, "成功", "文件处理完成��\n保存路径: " + outputPath);
    } else {
        QMessageBox::critical(this, "错误", "文件处理失败！");
    }
}

void MainWindow::selectUntranslatedSrt() {
    QString filePath = QFileDialog::getOpenFileName(
        this,
        "选择未翻译SRT文件",
        QString(),
        "SRT文件 (*.srt);;所有文件 (*.*)"
    );
    
    if (!filePath.isEmpty()) {
        untranslatedSrtPathEdit->setText(filePath);
    }
}

void MainWindow::selectTranslatedSrt() {
    QString filePath = QFileDialog::getOpenFileName(
        this,
        "选择已翻译SRT文件",
        QString(),
        "SRT文件 (*.srt);;所有文件 (*.*)"
    );
    
    if (!filePath.isEmpty()) {
        translatedSrtPathEdit->setText(filePath);
    }
}

void MainWindow::processSrt() {
    QString untranslatedPath = untranslatedSrtPathEdit->text();
    QString translatedPath = translatedSrtPathEdit->text();
    
    if (untranslatedPath.isEmpty() || translatedPath.isEmpty()) {
        QMessageBox::warning(this, "警告", "请先选择未翻译和已翻译的SRT文件");
        return;
    }

    // 构建输出路径
    QString outputDir = QDir::currentPath() + "/output";
    QString outputPath = outputDir + "/output.json";

    // 处理SRT文件
    if (srtProcessor.processSrtToJson(untranslatedPath, translatedPath, outputPath)) {
        QMessageBox::information(this, "成功", "文件处理完成！\n保存路径: " + outputPath);
    } else {
        QMessageBox::critical(this, "错误", "文件处理失败！");
    }
}

// 新增：获取关键词列表的方法
QStringList MainWindow::getKeywordsList(const QString& keywords) {
    QStringList result;
    if (keywords.isEmpty()) {
        return result;
    }
    
    // 分割关键词字符串
    result = keywords.split(",", Qt::SkipEmptyParts);
    
    // 去除每个关键词的首尾空格
    for (QString& keyword : result) {
        keyword = keyword.trimmed();
    }
    
    // 移除空字符串
    result.removeAll("");
    
    return result;
}

// 添加新的初始化关键词方法
void MainWindow::initializeKeywords() {
    // 加载JSON关键词
    QStringList savedJsonKeywords = loadKeywords(jsonKeywordsFileName);
    if (!savedJsonKeywords.isEmpty()) {
        QString keywordsText = savedJsonKeywords.join(",");
        jsonKeywordsEdit->setText(keywordsText);
        qDebug() << "已加载JSON关键词:" << keywordsText;
    }

    // 加载TXT关键词
    QStringList savedTxtKeywords = loadKeywords(txtKeywordsFileName);
    if (!savedTxtKeywords.isEmpty()) {
        QString keywordsText = savedTxtKeywords.join(",");
        txtKeywordsEdit->setText(keywordsText);
        qDebug() << "已加载TXT关键词:" << keywordsText;
    }
}

bool MainWindow::saveKeywords(const QStringList& keywords, const QString& fileName) {
    QString filePath = QDir::currentPath() + "/" + fileName;
    QFile file(filePath);
    
    if (!file.open(QIODevice::WriteOnly | QIODevice::Text)) {
        qDebug() << "无法创建关键词文件:" << filePath;
        return false;
    }
    
    QTextStream out(&file);
    for (const QString& keyword : keywords) {
        out << keyword << "\n";
    }
    
    file.close();
    qDebug() << "关键词已保存到:" << filePath;
    return true;
}

QStringList MainWindow::loadKeywords(const QString& fileName) {
    QStringList keywords;
    QString filePath = QDir::currentPath() + "/" + fileName;
    QFile file(filePath);
    
    if (!file.open(QIODevice::ReadOnly | QIODevice::Text)) {
        qDebug() << "无法打开关键词文件:" << filePath;
        return keywords;
    }
    
    QTextStream in(&file);
    while (!in.atEnd()) {
        QString line = in.readLine().trimmed();
        if (!line.isEmpty()) {
            keywords.append(line);
        }
    }
    
    file.close();
    return keywords;
}

void MainWindow::selectTxtToJsonInput() {
    QString filePath = QFileDialog::getOpenFileName(
        this,
        "选择TXT文件",
        QString(),
        "TXT文件 (*.txt);;所有文件 (*.*)"
    );
    
    if (!filePath.isEmpty()) {
        txtToJsonInputEdit->setText(filePath);
    }
}

void MainWindow::processTxtToJson() {
    QString inputPath = txtToJsonInputEdit->text();
    if (inputPath.isEmpty()) {
        QMessageBox::warning(this, "警告", "请先选择TXT文件");
        return;
    }

    // 构建输出路径
    QString outputDir = QDir::currentPath() + "/output";
    QString outputPath = outputDir + "/output.json";

    // 处理文件
    if (txtToJsonProcessor.processTxtToJson(inputPath, outputPath)) {
        QMessageBox::information(this, "成功", "文件处理完成！\n保存路径: " + outputPath);
    } else {
        QMessageBox::critical(this, "错误", "文件处理失败！");
    }
}

void MainWindow::selectJsonToSrtInput() {
    QString filePath = QFileDialog::getOpenFileName(
        this,
        "选择JSON文件",
        QString(),
        "JSON文件 (*.json);;所有文件 (*.*)"
    );
    
    if (!filePath.isEmpty()) {
        jsonToSrtInputEdit->setText(filePath);
    }
}

void MainWindow::processJsonToSrt() {
    QString inputPath = jsonToSrtInputEdit->text();
    if (inputPath.isEmpty()) {
        QMessageBox::warning(this, "警告", "请先选择JSON文件");
        return;
    }

    // 构建输出路径
    QString outputDir = QDir::currentPath() + "/output";
    QString sourceSrtPath = outputDir + "/source.srt";
    QString translatedSrtPath = outputDir + "/translated.srt";

    // 处理文件
    if (jsonToSrtProcessor.processJsonToSrt(inputPath, sourceSrtPath, translatedSrtPath)) {
        QMessageBox::information(this, "成功", 
            QString("文件处理完成！\n源语言SRT: %1\n翻译SRT: %2")
            .arg(sourceSrtPath)
            .arg(translatedSrtPath));
    } else {
        QMessageBox::critical(this, "错误", "文件处理失败！");
    }
}
